module recon_type_mod
  use smooth_indicators_mod
  implicit none
  
  type :: recon_type
    logical :: initialized = .false.
    integer :: id          = 0                       ! Sub-stencil ID
    integer :: nd          = 0                       ! Dimension number
    integer :: sw          = 0                       ! Stencil width for isotropy stencil
    integer :: swx         = 0                       ! Stencil width in x-dir
    integer :: swy         = 0                       ! Stencil width in y-dir
    integer :: swz         = 0                       ! Stencil width in z-dir
    integer :: sub_sw      = 0                       ! Sub-stencil width
    integer :: sub_nc      = 0                       ! Sub-stencil cell number
    integer :: nc          = 0                       ! Number of cells
    integer :: ac          = 0                       ! Number of available cells
    integer :: ns          = 0                       ! Number of sub-stencils
    integer :: nterm       = 0                       ! Number of polynomial terms of sub-stencil
    integer :: npt         = 0                       ! Number of evaluation points
    integer :: ic          = 0                       ! Index of center cell on x direction
    integer :: jc          = 0                       ! Index of center cell on y direction
    integer :: kc          = 0                       ! Index of center cell on z direction
    integer :: is          = 0                       ! Start index of subarray
    integer :: ie          = 0                       ! End index of subarray
    integer :: js          = 0                       ! Start index of subarray
    integer :: je          = 0                       ! End index of subarray
    integer :: ks          = 0                       ! Start index of subarray
    integer :: ke          = 0                       ! End index of subarray
    integer :: rw          = 0                       ! Reconstruction width
    integer :: rwL         = 0                       ! Reconstruction width on left side
    integer :: rwR         = 0                       ! Reconstruction width on right side
    integer :: rwB         = 0                       ! Reconstruction width on bottom side
    integer :: rwT         = 0                       ! Reconstruction width on top side
    integer :: rwD         = 0                       ! Reconstruction width on down side
    integer :: rwU         = 0                       ! Reconstruction width on up side
    integer :: max_ngb     = 0                       ! max neighbour cells on cell, for TCI
    real(16):: dx
    real(16):: dy
    real(16):: dz
    
    real(8)  :: eps     = 1.e-15
    real(16) :: eps_r16 = 1.e-17
    
    integer , allocatable :: cell_mask     (:,:,:)     ! Mask unavailable cells by 0
    integer , allocatable :: poly_mask     (:,:,:)     ! Mask unavailable polynomial terms by 0
    integer , allocatable :: ijk_to_1d     (:,:,:)     ! Map i,j index to 1D cell index. Just for rhombus stencil
    real(16), allocatable :: xc            (:    )     ! X coordinate of cell center
    real(16), allocatable :: yc            (:    )     ! Y coordinate of cell center
    real(16), allocatable :: zc            (:    )     ! Z coordinate of cell center
    real(16), allocatable :: x             (:    )     ! X coordinate of evaluation point
    real(16), allocatable :: y             (:    )     ! Y coordinate of evaluation point
    real(16), allocatable :: z             (:    )     ! Y coordinate of evaluation point
    real( 8), allocatable :: poly          (:,:  )     ! Polynomial terms on each evaluation point
    real(16), allocatable :: poly_r16      (:,:  )     ! Polynomial terms on each evaluation point
    real( 8), allocatable :: dpoly         (:,:,:,:,:) ! Derivative polynomial terms on each evaluation point,(npt,nc,0:max_x_deriv_order,0:max_y_deriv_order)
    real(16), allocatable :: dpoly_r16     (:,:,:,:,:) ! Derivative polynomial terms on each evaluation point,(npt,nc,0:max_x_deriv_order,0:max_y_deriv_order)
    real( 8), allocatable :: iA            (:,:  )     ! A * a = f, iA = inverse(A)
    real( 8), allocatable :: recon_mtx     (:,:  )     ! Reconstruction matrix, poly * iA
    real(16), allocatable :: recon_mtx_r16 (:,:  )     ! Reconstruction matrix, poly * iA
    real( 8), allocatable :: drecon_mtx    (:,:,:,:,:) ! Reconstruction derivative matrix, poly * iA, (npt,nc,0:max_x_deriv_order,0:max_y_deriv_order,0:max_z_deriv_order)
    real(16), allocatable :: drecon_mtx_r16(:,:,:,:,:) ! Reconstruction derivative matrix, poly * iA, (npt,nc,0:max_x_deriv_order,0:max_y_deriv_order,0:max_z_deriv_order)
    real( 8), allocatable :: gamma         (:,:  )     ! Ideal coefficients for combining sub-stencils (only on stencil)
    real( 8), allocatable :: rp            (:,:  )
    real( 8), allocatable :: rn            (:,:  )
    real( 8), allocatable :: sigmap        (:    )
    real( 8), allocatable :: sigman        (:    )
    integer , allocatable :: cell_ngb      (:,:,:,:,:) ! For trouble cell indicator, neighbor cells which are not masked. dim 1: directions, dim 2: number of neighbour cells, dim 3~5: i,j,k cell indices
    integer , allocatable :: cell_ngb_size (    :,:,:) ! For trouble cell indicator, number of neighbor cells. dim 1~3: i,j,k cell indices
    
    type(recon_type), allocatable :: subs(:) ! Sub-stencils
    
    integer, pointer :: ijxy(:,:) => null() ! Just for rhombus stencil
    
    procedure(smooth_indicator_interface            ), nopass, pointer :: smooth_indicator       => null()
    procedure(recon_init_interface                  ),   pass, pointer :: init                   => null()
    procedure(recon_add_point_interface             ),   pass, pointer :: add_point              => recon_add_point
    procedure(recon_calc_recon_matrix_interface     ),   pass, pointer :: calc_recon_matrix      => null()
    procedure(recon_calc_ideal_coefs_interface      ),   pass, pointer :: calc_ideal_coefs       => null()
    procedure(recon_calc_poly_interface             ),   pass, pointer :: calc_poly              => null()
    procedure(recon_calc_iA_interface               ),   pass, pointer :: calc_iA                => null()
    procedure(recon_reconstruct_interface           ),   pass, pointer :: reconstruct            => null()
    procedure(recon_reconstruct_by_poly_interface   ),   pass, pointer :: reconstruct_by_poly    => null()
    procedure(recon_reconstruct_deriv_interface     ),   pass, pointer :: reconstruct_deriv      => null()
    procedure(recon_trouble_cell_indicator_interface),   pass, pointer :: trouble_cell_indicator => null()
    procedure(recon_release_unused_memory_interface ),   pass, pointer :: release_unused_memory  => recon_release_unused_memory
    procedure(recon_clear_interface                 ),   pass, pointer :: clear                  => recon_clear
    
    procedure(recon_get_jacobian_inerface), pass, pointer :: get_jab => null()
    
  contains
    procedure, private :: recon_assign
    generic :: assignment(=) => recon_assign
    !final :: recon_final
  end type recon_type
  
  interface
    subroutine recon_init_interface(this, nd, sw, swx, swy, swz, xc, yc, zc, dx, dy, dz, ic, jc, kc, is, ie, js, je, ks, ke, mask, id)
      import recon_type
      class(recon_type), intent(inout) :: this
      integer , intent(in), optional :: nd
      integer , intent(in), optional :: sw
      integer , intent(in), optional :: swx
      integer , intent(in), optional :: swy
      integer , intent(in), optional :: swz
      real(16), intent(in), optional :: xc(:)
      real(16), intent(in), optional :: yc(:)
      real(16), intent(in), optional :: zc(:)
      real(16), intent(in), optional :: dx
      real(16), intent(in), optional :: dy
      real(16), intent(in), optional :: dz
      integer , intent(in), optional :: ic
      integer , intent(in), optional :: jc
      integer , intent(in), optional :: kc
      integer , intent(in), optional :: is
      integer , intent(in), optional :: ie
      integer , intent(in), optional :: js
      integer , intent(in), optional :: je
      integer , intent(in), optional :: ks
      integer , intent(in), optional :: ke
      integer , intent(in), optional :: mask(:,:,:) ! Cell mask
      integer , intent(in), optional :: id
    end subroutine recon_init_interface
    
    recursive subroutine recon_add_point_interface(this, x, y, z)
      import recon_type
      class(recon_type), intent(inout) :: this
      real(16), intent(in), optional :: x
      real(16), intent(in), optional :: y
      real(16), intent(in), optional :: z
    end subroutine recon_add_point_interface
    
    subroutine recon_calc_recon_matrix_interface(this, ierr)
      import recon_type
      class(recon_type), intent(inout) :: this
      integer, intent(out) :: ierr
    end subroutine recon_calc_recon_matrix_interface
    
    subroutine recon_calc_ideal_coefs_interface(this, ierr)
      import recon_type
      class(recon_type), intent(inout) :: this
      integer, intent(out) :: ierr
    end subroutine recon_calc_ideal_coefs_interface
    
    subroutine recon_calc_poly_interface(this, ierr)
      import recon_type
      class(recon_type), intent(inout) :: this
      integer, intent(out) :: ierr
    end subroutine recon_calc_poly_interface
    
    subroutine recon_calc_iA_interface(this, ierr)
      import recon_type
      class(recon_type), intent(inout) :: this
      integer, intent(out) :: ierr
    end subroutine recon_calc_iA_interface
    
    subroutine recon_reconstruct_interface(this, fi, fo, TCI, ierr)
      import recon_type
      class(recon_type), intent(inout) :: this
      real(8), intent(in ) :: fi(:,:,:) ! Cell averaged function values
      real(8), intent(out) :: fo(:)     ! Reconstructed function values on evaluation points
      logical, intent(in ), optional :: TCI ! Trouble Cell Indicator(TCI), 1 for existing TC(use WENO), 0 for no TC(use poly)
      integer, intent(out), optional :: ierr
    end subroutine recon_reconstruct_interface
    
    subroutine recon_reconstruct_by_poly_interface(this, poly, fi, fo, TCI, ierr)
      import recon_type
      class(recon_type), intent(inout) :: this
      real(8), intent(in ) :: poly(:,:) ! Polynomial with dimension of (npt,this%ac)
      real(8), intent(in ) :: fi(:,:,:) ! Cell averaged function values, total length must be equal to this%ac
      real(8), intent(out) :: fo(:)     ! Reconstructed function values on evaluation points, dimension (npt)
      logical, intent(in ), optional :: TCI ! Trouble Cell Indicator(TCI), 1 for existing TC(use WENO), 0 for no TC(use poly)
      integer, intent(out), optional :: ierr
    end subroutine recon_reconstruct_by_poly_interface
    
    subroutine recon_reconstruct_deriv_interface(this, fi, fx, fy, fz, ierr)
      import recon_type
      class(recon_type), intent(inout) :: this
      real(8), intent(in ) :: fi(:,:,:) ! Cell averaged function values
      real(8), intent(out), optional :: fx(:)     ! Reconstructed x derivative values on evaluation points
      real(8), intent(out), optional :: fy(:)     ! Reconstructed y derivative values on evaluation points
      real(8), intent(out), optional :: fz(:)     ! Reconstructed y derivative values on evaluation points
      integer, intent(out), optional :: ierr
    end subroutine recon_reconstruct_deriv_interface
    
    function recon_trouble_cell_indicator_interface(this, fi, dx) result(res)
      import recon_type
      class(recon_type), intent(in) :: this
      logical res
      real(8), intent(in) :: fi(:,:,:)
      real(8), intent(in) :: dx
    end function recon_trouble_cell_indicator_interface
    
    subroutine recon_release_unused_memory_interface(this)
      import recon_type
      class(recon_type), intent(inout) :: this
    end subroutine recon_release_unused_memory_interface
    
    subroutine recon_clear_interface(this)
      import recon_type
      class(recon_type), intent(inout) :: this
    end subroutine recon_clear_interface
    
    subroutine recon_assign_interface(this, other)
      import recon_type
      class(recon_type), intent(inout) :: this
      class(recon_type), intent(in   ) :: other
    end subroutine recon_assign_interface
    
    subroutine recon_get_jacobian_inerface(this, dqLdq, dqRdq, q )
      import recon_type
      class(recon_type), intent(in   )           :: this
      real(8)          , intent(inout)           :: dqLdq(this%sw)
      real(8)          , intent(inout)           :: dqRdq(this%sw)
      real(8)          , intent(in   ), optional :: q    (this%sw)
    end subroutine recon_get_jacobian_inerface
  end interface
  
  contains

  subroutine recon_add_point(this, x, y, z)

    class(recon_type), intent(inout) :: this
    real(16), intent(in), optional :: x
    real(16), intent(in), optional :: y
    real(16), intent(in), optional :: z

    real(16), allocatable :: tmp(:)
    integer i, iStencil

    this%npt = this%npt + 1

    ! x-dir
    if (present(x)) then
      allocate(tmp(this%npt))
      do i = 1, this%npt - 1
        tmp(i) = this%x(i)
      end do
      tmp(this%npt) = x
      if (allocated(this%x)) deallocate(this%x)
      allocate(this%x(this%npt))
      this%x = tmp
    endif
    
    ! y-dir
    if (present(y)) then
      do i = 1, this%npt - 1
        tmp(i) = this%y(i)
      end do
      tmp(this%npt) = y
      if (allocated(this%y)) deallocate(this%y)
      allocate(this%y(this%npt))
      this%y = tmp
    end if
    
    ! z-dir
    if (present(z)) then
      do i = 1, this%npt - 1
        tmp(i) = this%z(i)
      end do
      tmp(this%npt) = z
      if (allocated(this%z)) deallocate(this%z)
      allocate(this%z(this%npt))
      this%z = tmp
    end if
    
    deallocate(tmp)

    if (allocated(this%subs)) then
      ! Add point to sub-stencils.
      do iStencil = 1, this%ns
        if(this%nd==1)call this%subs(iStencil)%add_point(x)
        if(this%nd==2)call this%subs(iStencil)%add_point(x, y)
        if(this%nd==3)call this%subs(iStencil)%add_point(x, y, z)
      end do
    end if

  end subroutine recon_add_point

  subroutine recon_release_unused_memory(this)

    class(recon_type), intent(inout) :: this

    type(recon_type), allocatable :: subs(:)
    real(8), allocatable :: gamma(:,:)
    integer i, k, ns

    if (allocated(this%poly_mask     )) deallocate(this%poly_mask     )
    if (allocated(this%ijk_to_1d     )) deallocate(this%ijk_to_1d     )
    if (allocated(this%xc            )) deallocate(this%xc            )
    if (allocated(this%yc            )) deallocate(this%yc            )
    if (allocated(this%x             )) deallocate(this%x             )
    if (allocated(this%y             )) deallocate(this%y             )
    if (allocated(this%poly_r16      )) deallocate(this%poly_r16      )
    if (allocated(this%dpoly_r16     )) deallocate(this%dpoly_r16     )
    if (allocated(this%recon_mtx_r16 )) deallocate(this%recon_mtx_r16 )
    if (allocated(this%drecon_mtx_r16)) deallocate(this%drecon_mtx_r16)

    ! Shrink subs and gamma arrays to only contain unmasked sub-stencils.
    ns = this%ns
    do k = 1, this%ns
      if (any(this%subs(k)%cell_mask == 0)) ns = ns - 1
    end do
    if (ns < this%ns) then
      allocate(subs(ns), gamma(ns,this%npt))
      i = 1
      do k = 1, this%ns
        if (any(this%subs(k)%cell_mask == 0)) cycle
        subs(i) = this%subs(k)
        gamma(i,:) = this%gamma(k,:)
        i = i + 1
      end do
      deallocate(this%subs, this%gamma)
      this%ns = ns
      allocate(this%subs(ns), this%gamma(ns,this%npt))
      do k = 1, this%ns
        this%subs(k) = subs(k)
        this%gamma(k,:) = gamma(k,:)
      end do
      deallocate(subs, gamma)
    end if

  end subroutine recon_release_unused_memory
    
  subroutine recon_clear(this)

    class(recon_type), intent(inout) :: this
    
    this%initialized = .false.
    this%id          = 0
    this%nd          = 0
    this%sw          = 0
    this%swx         = 0
    this%swy         = 0
    this%swz         = 0
    this%sub_sw      = 0
    this%sub_nc      = 0
    this%nc          = 0
    this%ns          = 0
    this%nterm       = 0
    this%npt         = 0
    this%ic          = 0
    this%jc          = 0
    this%kc          = 0
    this%is          = 0
    this%ie          = 0
    this%js          = 0
    this%je          = 0
    this%ks          = 0
    this%ke          = 0
    this%max_ngb     = 0
    
    if (allocated(this%cell_mask     )) deallocate(this%cell_mask     )
    if (allocated(this%poly_mask     )) deallocate(this%poly_mask     )
    if (allocated(this%ijk_to_1d     )) deallocate(this%ijk_to_1d     )
    if (allocated(this%xc            )) deallocate(this%xc            )
    if (allocated(this%yc            )) deallocate(this%yc            )
    if (allocated(this%zc            )) deallocate(this%zc            )
    if (allocated(this%x             )) deallocate(this%x             )
    if (allocated(this%y             )) deallocate(this%y             )
    if (allocated(this%z             )) deallocate(this%z             )
    if (allocated(this%poly          )) deallocate(this%poly          )
    if (allocated(this%poly_r16      )) deallocate(this%poly_r16      )
    if (allocated(this%dpoly         )) deallocate(this%dpoly         )
    if (allocated(this%dpoly_r16     )) deallocate(this%dpoly_r16     )
    if (allocated(this%iA            )) deallocate(this%iA            )
    if (allocated(this%recon_mtx     )) deallocate(this%recon_mtx     )
    if (allocated(this%recon_mtx_r16 )) deallocate(this%recon_mtx_r16 )
    if (allocated(this%drecon_mtx    )) deallocate(this%drecon_mtx    )
    if (allocated(this%drecon_mtx_r16)) deallocate(this%drecon_mtx_r16)
    if (allocated(this%gamma         )) deallocate(this%gamma         )
    if (allocated(this%rp            )) deallocate(this%rp            )
    if (allocated(this%rn            )) deallocate(this%rn            )
    if (allocated(this%sigmap        )) deallocate(this%sigmap        )
    if (allocated(this%sigman        )) deallocate(this%sigman        )
    if (allocated(this%cell_ngb      )) deallocate(this%cell_ngb      )
    if (allocated(this%cell_ngb_size )) deallocate(this%cell_ngb_size )
    if (allocated(this%subs          )) deallocate(this%subs          )
    
    this%smooth_indicator => null()
    
  end subroutine recon_clear

  subroutine recon_assign(this, other)

    class(recon_type), intent(inout) :: this
    class(recon_type), intent(in   ) :: other

    call this%clear()

    this%initialized = .true.
    this%id          = other%id         
    this%nd          = other%nd         
    this%sw          = other%sw         
    this%swx         = other%swx        
    this%swy         = other%swy         
    this%swz         = other%swz       
    this%sub_sw      = other%sub_sw     
    this%sub_nc      = other%sub_nc     
    this%nc          = other%nc         
    this%ns          = other%ns         
    this%nterm       = other%nterm      
    this%npt         = other%npt        
    this%ic          = other%ic         
    this%jc          = other%jc         
    this%kc          = other%kc         
    this%is          = other%is         
    this%ie          = other%ie         
    this%js          = other%js         
    this%je          = other%je      
    this%ks          = other%ks         
    this%ke          = other%ke      
    this%max_ngb     = other%max_ngb
    
    if(allocated(other%cell_mask     )) allocate(this%cell_mask     (this%swx,this%swy,this%swz))
    if(allocated(other%poly_mask     )) allocate(this%poly_mask     (this%swx,this%swy,this%swz))
    if(allocated(other%ijk_to_1d     )) allocate(this%ijk_to_1d     (this%swx,this%swy,this%swz))
    if(allocated(other%xc            )) allocate(this%xc            (this%swx))
    if(allocated(other%yc            )) allocate(this%yc            (this%swy))
    if(allocated(other%zc            )) allocate(this%zc            (this%swz))
    if(allocated(other%x             )) allocate(this%x             (this%npt))
    if(allocated(other%y             )) allocate(this%y             (this%npt))
    if(allocated(other%z             )) allocate(this%z             (this%npt))
    if(allocated(other%poly          )) allocate(this%poly          (this%npt,this%nc))
    if(allocated(other%poly_r16      )) allocate(this%poly_r16      (this%npt,this%nc))
    if(allocated(other%dpoly         )) allocate(this%dpoly         (this%npt,this%nc,0:this%swx-1,0:this%swy-1,0:this%swz-1))
    if(allocated(other%dpoly_r16     )) allocate(this%dpoly_r16     (this%npt,this%nc,0:this%swx-1,0:this%swy-1,0:this%swz-1))
    if(allocated(other%iA            )) allocate(this%iA            (this%nc ,this%nc))
    if(allocated(other%recon_mtx     )) allocate(this%recon_mtx     (this%npt,this%nc))
    if(allocated(other%recon_mtx_r16 )) allocate(this%recon_mtx_r16 (this%npt,this%nc))
    if(allocated(other%drecon_mtx    )) allocate(this%drecon_mtx    (this%npt,this%nc,0:this%swx-1,0:this%swy-1,0:this%swz-1))
    if(allocated(other%drecon_mtx_r16)) allocate(this%drecon_mtx_r16(this%npt,this%nc,0:this%swx-1,0:this%swy-1,0:this%swz-1))
    if(allocated(other%gamma         )) allocate(this%gamma         (this%ns,this%npt))
    if(allocated(other%rp            )) allocate(this%rp            (this%ns,this%npt))
    if(allocated(other%rn            )) allocate(this%rn            (this%ns,this%npt))
    if(allocated(other%sigmap        )) allocate(this%sigmap        (this%npt))
    if(allocated(other%sigman        )) allocate(this%sigman        (this%npt))
    if(allocated(other%cell_ngb      )) allocate(this%cell_ngb      (this%nd,this%max_ngb,this%swx,this%swy,this%swz))
    if(allocated(other%cell_ngb_size )) allocate(this%cell_ngb_size (                     this%swx,this%swy,this%swz))
    if(allocated(other%subs          )) allocate(this%subs          (this%ns))

    this%cell_mask = other%cell_mask
    this%poly      = other%poly
    this%iA        = other%iA
    this%recon_mtx = other%recon_mtx

    this%ijxy                   => other%ijxy
    this%smooth_indicator       => other%smooth_indicator            
    this%init                   => other%init                  
    this%calc_recon_matrix      => other%calc_recon_matrix     
    this%calc_ideal_coefs       => other%calc_ideal_coefs      
    this%reconstruct            => other%reconstruct           
    this%trouble_cell_indicator => other%trouble_cell_indicator

  end subroutine recon_assign
  
  subroutine recon_final(this)

    type(recon_type), intent(inout) :: this

    call this%clear()

  end subroutine recon_final
end module recon_type_mod
